import torch
import numpy as np

def dwconv3x3_transform(weight):
    weight_t = torch.zeros(0)
    for i in range(int(weight.size(0)/32)):
        temp = weight[i*32:(i+1)*32].permute(1,2,3,0)
        weight_t = torch.cat([weight_t, temp])
    weight_t = weight_t.reshape(-1,32)
    return weight_t.detach().numpy().flatten()

def pwconv1x1_transform(weight):
    oc = weight.size(0)
    ic = weight.size(1)
    weight_np = weight.detach().numpy().flatten()
    weight_t = np.zeros(oc*ic, dtype = np.float32)
    for Mx in range(int(oc/32)):
        for Nx in range(int(ic/32)):
            for n in range(32):
                for m in range(32):
                    weight_t[int((Mx*ic/32+Nx)*1024+n*32+m)] = weight_np[(Mx*32+m)*ic+Nx*32+n]
    return weight_t

net = torch.load('./EOQSkyNet.pth')
conv_layers = []
bn_layers = []
for module in net.named_modules():
    if isinstance(module[1], torch.nn.Conv2d):
        conv_layers.append(module[1])
    if isinstance(module[1], torch.nn.BatchNorm2d):
        bn_layers.append(module[1])

weight = np.zeros(0, dtype = np.int8)
for i in range(len(conv_layers)):
    if(conv_layers[i].groups>1):
        w = dwconv3x3_transform(conv_layers[i].weight)
    else:
        w = pwconv1x1_transform(conv_layers[i].weight)
    weight = np.concatenate((weight,w.astype(np.int8)))
weight.tofile('weight/SkyNetT.wt')

biasm = np.zeros(0, dtype = np.int32)
for i in range(len(conv_layers)):
    b = conv_layers[i].bias.detach().cpu().numpy().astype(np.int32)
    m = bn_layers[i].weight.detach().cpu().numpy().astype(np.int32)
    biasm = np.concatenate((biasm,b))
    biasm = np.concatenate((biasm,m))
biasm.tofile('weight/SkyNetT.bm')
print(weight.shape, biasm.shape)